package edu.nepu.flink.api.state;

import org.apache.flink.api.common.functions.RichMapFunction;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;

/**
 * @Date 2024/3/2 10:42
 * @Created by chenshuaijun
 */
public class OperatorUnionListStateDemo {

    public static void main(String[] args) throws Exception {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);
        DataStreamSource<String> streamSource = env.socketTextStream("hadoop102", 9999);
        /**
         * 下面我将演示算子状态的使用，想要使用算子状态我们需要实现CheckpointedFunction接口.
         * 我们的需求是在map算子中计算流过该算子的数据的条数
         */
        streamSource.map(new MyState()).print();


        env.execute();
    }

    static class MyState extends RichMapFunction<String, Integer> implements CheckpointedFunction{
        Integer number;
        ListState<Integer> listState;
        @Override
        public void open(Configuration parameters) throws Exception {
            super.open(parameters);
            number = 0;
        }

        /**
         * 将本地变量存储在算子状态中，这个是固定的写法
         * @param value The input value.
         * @return
         * @throws Exception
         */
        @Override
        public Integer map(String value) throws Exception {
            return ++number;
        }

        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {
            listState.clear();
            listState.add(number);
        }

        /**
         * 当从checkpoint通恢复的时候，会首先调用这个方法，初始化状态
         * @param context the context for initializing the operator
         * @throws Exception
         */
        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {
            // TODO 和ListState唯一的不同就是在这里
            listState = context.getOperatorStateStore().getUnionListState(new ListStateDescriptor<Integer>("listState",Types.INT));
            Iterable<Integer> integers = listState.get();
            // 这个是判断 状态是否恢复了
            if (context.isRestored()) {
                for (Integer value : integers) {
                    number += value;
                }
            }
        }
    }

}
